package main

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"k8s.io/client-go/rest"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
)

// === 数据结构保持不变 ===
type ThreadPoolBean struct {
	ActiveCount        int    `json:"activeCount"`
	PoolSize           int    `json:"poolSize"`
	CorePoolSize       int    `json:"corePoolSize"`
	MaxPoolSize        int    `json:"maxPoolSize"`
	KeepAliveSeconds   int    `json:"keepAliveSeconds"`
	ThreadPoolBeanName string `json:"threadPoolBeanName"`
}

type DataSourceStatus struct {
	IdleConnections           int `json:"idleConnections"`
	ActiveConnections         int `json:"activeConnections"`
	TotalConnections          int `json:"totalConnections"`
	ThreadsAwaitingConnection int `json:"threadsAwaitingConnection"`
	MaximumPoolSize           int `json:"maximumPoolSize"`
}

type MonitorResult struct {
	ThreadPoolStatus []ThreadPoolBean `json:"threadPoolStatus"`
	DataSourceStatus DataSourceStatus `json:"dataSourceStatus"`
}

type ApiResponse struct {
	Status int           `json:"status"`
	Msg    string        `json:"msg"`
	Result MonitorResult `json:"result"`
}

// === Prometheus 指标（保持不变）===
var (
	threadPoolActiveCount = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{Name: "thread_pool_active_threads", Help: "Number of active threads in the thread pool"},
		[]string{"service", "pod", "namespace", "thread_pool"},
	)
	threadPoolPoolSize = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{Name: "thread_pool_current_size", Help: "Current number of threads in the thread pool"},
		[]string{"service", "pod", "namespace", "thread_pool"},
	)
	threadPoolCoreSize = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{Name: "thread_pool_core_size", Help: "Core number of threads in the thread pool"},
		[]string{"service", "pod", "namespace", "thread_pool"},
	)
	threadPoolMaxSize = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{Name: "thread_pool_max_size", Help: "Maximum number of threads in the thread pool"},
		[]string{"service", "pod", "namespace", "thread_pool"},
	)
	threadPoolKeepAliveSeconds = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{Name: "thread_pool_keep_alive_seconds", Help: "Keep-alive time for threads in the thread pool"},
		[]string{"service", "pod", "namespace", "thread_pool"},
	)

	dataSourceIdleConnections = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{Name: "datasource_idle_connections", Help: "Number of idle connections in the datasource pool"},
		[]string{"service", "pod", "namespace"},
	)
	dataSourceActiveConnections = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{Name: "datasource_active_connections", Help: "Number of active connections in the datasource pool"},
		[]string{"service", "pod", "namespace"},
	)
	dataSourceTotalConnections = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{Name: "datasource_total_connections", Help: "Total number of connections in the datasource pool"},
		[]string{"service", "pod", "namespace"},
	)
	dataSourceAwaitingConnection = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{Name: "datasource_threads_awaiting_connection", Help: "Number of threads waiting for a datasource connection"},
		[]string{"service", "pod", "namespace"},
	)
	dataSourceMaxPoolSize = prometheus.NewGaugeVec(
		prometheus.GaugeOpts{Name: "datasource_max_pool_size", Help: "Maximum size of the datasource connection pool"},
		[]string{"service", "pod", "namespace"},
	)
)

func init() {
	prometheus.MustRegister(threadPoolActiveCount, threadPoolPoolSize, threadPoolCoreSize, threadPoolMaxSize, threadPoolKeepAliveSeconds,
		dataSourceIdleConnections, dataSourceActiveConnections, dataSourceTotalConnections, dataSourceAwaitingConnection, dataSourceMaxPoolSize)
}

// === 活跃指标追踪 ===
var (
	mu                   sync.RWMutex
	activePodKeys        = make(map[string]struct{}) // "namespace/service/pod"
	activeThreadPoolKeys = make(map[string]struct{}) // "namespace/service/pod/threadpool"
)

// getLabelKey 生成唯一标识
func getLabelKey(namespace, service, pod string) string {
	return fmt.Sprintf("%s/%s/%s", namespace, service, pod)
}

func getThreadPoolLabelKey(namespace, service, pod, threadPool string) string {
	return fmt.Sprintf("%s/%s/%s/%s", namespace, service, pod, threadPool)
}

// cleanupStaleMetrics 删除已不存在的 Pod 和线程池的指标
func cleanupStaleMetrics(currentPodKeys, currentThreadPoolKeys map[string]struct{}) {
	mu.Lock()
	defer mu.Unlock()

	// 清理不再存在的 Pod 指标
	for key := range activePodKeys {
		if _, exists := currentPodKeys[key]; !exists {
			split := strings.Split(key, "/")
			if len(split) != 3 {
				continue
			}
			labels := prometheus.Labels{
				"namespace": split[0],
				"service":   split[1],
				"pod":       split[2],
			}
			dataSourceIdleConnections.Delete(labels)
			dataSourceActiveConnections.Delete(labels)
			dataSourceTotalConnections.Delete(labels)
			dataSourceAwaitingConnection.Delete(labels)
			dataSourceMaxPoolSize.Delete(labels)
			log.Printf("🧹 Cleaned up metrics for pod: %s", key)
		}
	}

	// 清理不再存在的线程池指标
	for key := range activeThreadPoolKeys {
		if _, exists := currentThreadPoolKeys[key]; !exists {
			split := strings.Split(key, "/")
			if len(split) != 4 {
				continue
			}
			labels := prometheus.Labels{
				"namespace":   split[0],
				"service":     split[1],
				"pod":         split[2],
				"thread_pool": split[3],
			}
			threadPoolActiveCount.Delete(labels)
			threadPoolPoolSize.Delete(labels)
			threadPoolCoreSize.Delete(labels)
			threadPoolMaxSize.Delete(labels)
			threadPoolKeepAliveSeconds.Delete(labels)
			log.Printf("🧹 Cleaned up metrics for threadPool: %s", key)
		}
	}

	// 更新全局活跃状态
	activePodKeys = currentPodKeys
	activeThreadPoolKeys = currentThreadPoolKeys
}

// fetchPodMetrics 获取单个 Pod 的指标
func fetchPodMetrics(client *http.Client, pod v1.Pod, svc TargetService, currentPodKeys, currentThreadPoolKeys map[string]struct{}) {
	podIP := pod.Status.PodIP
	port := svc.Port
	url := fmt.Sprintf("http://%s:%d%s", podIP, port, svc.Path)

	resp, err := client.Get(url)
	if err != nil {
		log.Printf("❌ Error fetching %s/%s: %v", svc.Name, pod.Name, err)
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		body, _ := io.ReadAll(resp.Body)
		log.Printf("❌ HTTP %d from %s/%s: %s", resp.StatusCode, svc.Name, pod.Name, string(body))
		return
	}

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		log.Printf("❌ Error reading body from %s/%s: %v", svc.Name, pod.Name, err)
		return
	}

	var apiResp ApiResponse
	err = json.Unmarshal(body, &apiResp)
	if err != nil {
		log.Printf("❌ Error unmarshaling JSON from %s/%s: %v", svc.Name, pod.Name, err)
		return
	}

	commonLabels := prometheus.Labels{
		"service":   svc.Name,
		"pod":       pod.Name,
		"namespace": svc.Namespace,
	}

	// 更新数据源指标
	ds := apiResp.Result.DataSourceStatus
	dataSourceIdleConnections.With(commonLabels).Set(float64(ds.IdleConnections))
	dataSourceActiveConnections.With(commonLabels).Set(float64(ds.ActiveConnections))
	dataSourceTotalConnections.With(commonLabels).Set(float64(ds.TotalConnections))
	dataSourceAwaitingConnection.With(commonLabels).Set(float64(ds.ThreadsAwaitingConnection))
	dataSourceMaxPoolSize.With(commonLabels).Set(float64(ds.MaximumPoolSize))

	// 更新线程池指标
	for _, tp := range apiResp.Result.ThreadPoolStatus {
		tpLabels := prometheus.Labels{
			"service":     svc.Name,
			"pod":         pod.Name,
			"namespace":   svc.Namespace,
			"thread_pool": tp.ThreadPoolBeanName,
		}

		threadPoolActiveCount.With(tpLabels).Set(float64(tp.ActiveCount))
		threadPoolPoolSize.With(tpLabels).Set(float64(tp.PoolSize))
		threadPoolCoreSize.With(tpLabels).Set(float64(tp.CorePoolSize))
		threadPoolMaxSize.With(tpLabels).Set(float64(tp.MaxPoolSize))
		threadPoolKeepAliveSeconds.With(tpLabels).Set(float64(tp.KeepAliveSeconds))

		// 记录活跃线程池
		tpKey := getThreadPoolLabelKey(svc.Namespace, svc.Name, pod.Name, tp.ThreadPoolBeanName)
		currentThreadPoolKeys[tpKey] = struct{}{}
	}

	// 记录活跃 Pod
	podKey := getLabelKey(svc.Namespace, svc.Name, pod.Name)
	currentPodKeys[podKey] = struct{}{}

	log.Printf("✅ Updated metrics for %s/%s, %d thread pools", svc.Name, pod.Name, len(apiResp.Result.ThreadPoolStatus))
}

// startScraping 启动指标采集
func startScraping(kubeClient *kubernetes.Clientset, services []TargetService) {
	httpClient := &http.Client{
		Timeout: 2 * time.Second,
		Transport: &http.Transport{
			DisableKeepAlives: true, // 避免连接堆积
		},
	}
	ticker := time.NewTicker(15 * time.Second)

	go func() {
		for range ticker.C {
			currentPodKeys := make(map[string]struct{})
			currentThreadPoolKeys := make(map[string]struct{})

			for _, svc := range services {
				pods, err := kubeClient.CoreV1().Pods(svc.Namespace).List(context.TODO(), metav1.ListOptions{
					LabelSelector: fmt.Sprintf("app=%s", svc.Name),
				})
				if err != nil {
					log.Printf("❌ Error listing pods for %s: %v", svc.Name, err)
					continue
				}

				// 同步处理每个 Pod，避免并发写 map
				for _, pod := range pods.Items {
					if pod.Status.Phase == v1.PodRunning {
						fetchPodMetrics(httpClient, pod, svc, currentPodKeys, currentThreadPoolKeys)
					}
				}
			}

			// 所有 Pod 处理完成后，清理过期指标
			cleanupStaleMetrics(currentPodKeys, currentThreadPoolKeys)
		}
	}()
}

// TargetService 服务配置
type TargetService struct {
	Name      string
	Namespace string
	Port      int32
	Path      string
}

func main() {
	NS := os.Getenv("MY_POD_NAMESPACE")
	if len(NS) == 0 {
		NS = "default"
	}

	services := []TargetService{
		{Name: "school-dzb-dothandler", Namespace: NS, Port: 80, Path: "/dzb-dothandler/monitor/all/status"},
		{Name: "school-dzb-dothandler2", Namespace: NS, Port: 80, Path: "/dzb-dothandler/monitor/all/status"},
		{Name: "school-dzb-dothandler3", Namespace: NS, Port: 80, Path: "/dzb-dothandler/monitor/all/status"},
	}

	// 创建 Kubernetes 客户端
	config, err := rest.InClusterConfig()
	if err != nil {
		log.Fatalf("❌ Failed to get Kubernetes config: %v", err)
	}

	// use the current context in kubeconfig
	//config, err := clientcmd.BuildConfigFromFlags("", "D:\\code\\backend\\golangProject\\ql-ops-go\\conf\\kubeconfig\\dev")
	//config, err := clientcmd.BuildConfigFromFlags("", "/root/.kube/config")
	//if err != nil {
	//	panic(err.Error())
	//}

	clientset, err := kubernetes.NewForConfig(config)
	if err != nil {
		log.Fatalf("❌ Failed to create Kubernetes client: %v", err)
	}

	startScraping(clientset, services)

	// 注册 HTTP 路由
	http.Handle("/metrics", promhttp.Handler())
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte("Advanced ThreadPool & DataSource Exporter is running. Visit /metrics for metrics."))
	})

	log.Println("✅ Exporter started at :9090")
	log.Fatal(http.ListenAndServe(":9090", nil))
}
